"""
  描述：自行车模型（用于换道的车辆）
"""

import math
import numpy as np

from kalman_filter._02_ekf_bicyle_model import BicycleModelKalmanFilter

class BicycleModel:
  def __init__(self, dt, wheelbase=1, y0=0, x0=0, v0=0, psi0=0, process_noise_std=[0.3, 0.3, 0.05, 0.02], measurement_noise_std=[0.2, 0.8, 0.5, 0.1]):
    self.dt = dt  # 时间步长
    self.L = wheelbase  # 车辆轴距
    self.process_noise_std = process_noise_std
    self.measurement_noise_std = measurement_noise_std

    # 初始状态
    self.y = y0  # 纵向位置y
    self.x = x0  # 横向位置x
    self.psi = psi0  # 航向角psi（与道路垂直方向，即x轴的夹角）
    self.v = v0  # 轴向速度v
    self.v_y = v0 * math.sin(self.psi)  # 纵向速度v_y
    self.v_x = v0 * math.cos(self.psi)  # 横向速度v_x

    # 测量状态
    self.measure_y = y0
    self.measure_x = x0
    self.measure_psi = psi0
    self.measure_v = v0
    self.measure_v_y = v0 * math.sin(psi0)
    self.measure_v_x = v0 * math.cos(psi0)

    # 滤波状态
    self.kf = BicycleModelKalmanFilter(dt, wheelbase, process_noise_std, measurement_noise_std)
    self.kf.reset(np.array([x0, y0, v0, psi0])) # 初始状态估计
    self.filter_x, self.filter_y, self.filter_v, self.filter_psi = self.kf.get_state()
    self.filter_v_y = self.filter_v * math.sin(self.filter_psi)
    self.filter_v_x = self.filter_v * math.cos(self.filter_psi)
    P = self.kf.get_covariance()
    self.filter_covariances = [P[0,0], P[1,1], P[2,2], P[3,3]] # 位置、速度、航向角的协方差

  def _add_noise(self, data, stds_dev, noise_type="gaussian"):
    """添加噪声到数据"""
    stds_dev = np.array(stds_dev)
    if noise_type == "gaussian":
        # 高斯噪声
        noise = np.random.normal(0, stds_dev, data.shape)
    elif noise_type == "uniform":
        # 均匀噪声范围为 [-std, std]
        noise = np.random.uniform(-stds_dev, stds_dev, data.shape)
    return data + noise

  def update(self, a, delta):
    """
      描述：更新状态
      参数：
        a:  加速度
        delta: 转向角
    """
    # 更新真实状态
    self.y += self.v_y * self.dt + 1 / 2 * a * math.sin(self.psi) * (self.dt**2) # 孙康师兄PartC动力学模型
    self.x += self.v_x * self.dt + 1 / 2 * a * math.cos(self.psi) * (self.dt**2)
    # self.psi -= self.v * math.tan(delta) * self.dt # ?这里应该是减号吧，因为含义是和x轴的夹角？？？
    self.v += a * self.dt # todo 需要防止出现负速度？
    self.psi += (self.v / self.L) * math.tan(delta) * self.dt
    self.v_y = self.v * math.sin(self.psi)
    self.v_x = self.v * math.cos(self.psi)

    # 生成测量状态
    measurement_data = self._add_noise(np.array([self.x, self.y, self.v, self.psi]), self.measurement_noise_std) # 真实值增加噪声，形成观测值
    self.measure_x, self.measure_y, self.measure_v, self.measure_psi = measurement_data
    self.measure_v_y = self.measure_v * math.sin(self.measure_psi)
    self.measure_v_x = self.measure_v * math.cos(self.measure_psi)

    # 生成滤波状态
    self.kf.filter_step(measurement_data, [a, delta]) # 执行
    self.filter_x, self.filter_y, self.filter_v, self.filter_psi = self.kf.get_state()
    self.filter_v_y = self.filter_v * math.sin(self.filter_psi)
    self.filter_v_x = self.filter_v * math.cos(self.filter_psi)
    P = self.kf.get_covariance()
    self.filter_covariances = [P[0,0], P[1,1], P[2,2], P[3,3]] # 位置、速度、航向角的协方差

  def get_state(self):
    """
      描述：获取状态
    """
    return self.y, self.x, self.psi, self.v, self.v_y, self.v_x
  
  def get_measure_state(self):
    """
      描述：获取测量状态
    """
    return self.measure_y, self.measure_x, self.measure_psi, self.measure_v, self.measure_v_y, self.measure_v_x
  
  def get_filter_state(self):
    """
      描述：获取滤波状态
    """
    return self.filter_y, self.filter_x, self.filter_psi, self.filter_v, self.filter_v_y, self.filter_v_x, self.filter_covariances
  
  def calc_predict_state(self, a_arr, delta_arr):
    """
      描述：计算预测状态
      参数：
        a_arr:  加速度数组
        delta_arr: 转向角数组
      返回：
        预测状态数组
    """
    states = []
    y_pred = self.filter_y
    x_pred = self.filter_x
    psi_pred = self.filter_psi
    v_pred = self.filter_v
    v_y_pred = self.filter_v_y
    v_x_pred = self.filter_v_x
    for a, delta in zip(a_arr, delta_arr): # 用zip()函数同时遍历两个数组
      y_pred += v_y_pred * self.dt + 1 / 2 * a * math.sin(psi_pred) * (self.dt**2)
      x_pred += v_x_pred * self.dt + 1 / 2 * a * math.cos(psi_pred) * (self.dt**2)
      # psi_pred -= v_pred * math.tan(delta) * self.dt
      v_pred += a * self.dt
      psi_pred += (v_pred / self.L) * math.tan(delta) * self.dt
      v_y_pred = v_pred * math.sin(psi_pred)
      v_x_pred = v_pred * math.cos(psi_pred)
      states.append((y_pred, x_pred, psi_pred, v_pred, v_y_pred, v_x_pred))
    return np.array(states)
